#!/usr/bin/python
# encoding:utf-8

import click
from mlflow.utils.file_utils import TempDir
import mlflow
import waveletai
from hydraulic_sys_classification_model import fit_classifier


@click.command()
@click.option("--fault_type", "-ft", type=str, default='冷却器条件', help="故障类型")
@click.option("--conv_num", "-con", type=int, default=32, help="卷积核")
@click.option("--batch_size", "-bs", type=int, default=20, help="随机优化的minibatches的大小")
@click.option("--nb_epochs", "-ep", type=int, default=30, help="")
@click.option("--training_data", "-td", type=str, default='', help="默认数据集")
def train(fault_type, conv_num, batch_size, nb_epochs, training_data):
    return train_model(fault_type, conv_num, batch_size, nb_epochs, training_data)


def train_model(fault_type, conv_num, batch_size, nb_epochs, training_data):
    with mlflow.start_run():
        with TempDir() as tmp:
            waveletai.log_param("fault_type_param", str(fault_type))
            waveletai.log_param("conv_num_param", str(conv_num))
            waveletai.log_param("batch_size_param", str(batch_size))
            waveletai.log_param("nb_epochs_param", str(nb_epochs))
            waveletai.log_param("training_data_param", str(training_data))
            tmp_path = tmp.path()
            output_directory = tmp_path + '/'
            waveletai.init()
            waveletai.download_dataset_artifacts(training_data, output_directory,unzip=True)
            training_data_path = output_directory + "dataset.csv"
            fit_classifier(training_data_path, fault_type, conv_num, batch_size, nb_epochs, output_directory)


if __name__ == '__main__':
    train()
